#!/bin/bash

# Hostname Modification Script for 5GC Testbed VMs
# This script modifies the actual hostnames of each VM

set -e

# VM Configuration
VM1="192.168.1.206"  # Core Network
VM2="192.168.1.219"  # UPF
VM3="192.168.1.221"  # gNodeB
VM4="192.168.1.205"  # UE

# Target hostnames
HOSTNAME_VM1="vm1-5gc-core"
HOSTNAME_VM2="vm2-5gc-upf"
HOSTNAME_VM3="vm3-5gc-gnb"
HOSTNAME_VM4="vm4-5gc-ue"

# Colors for output
GREEN='\033[0;32m'
BLUE='\033[0;34m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

info() {
    echo -e "${BLUE}$1${NC}"
}

warn() {
    echo -e "${YELLOW}WARNING: $1${NC}"
}

error() {
    echo -e "${RED}ERROR: $1${NC}"
}

# Function to get current hostname of a VM
get_current_hostname() {
    local vm_ip=$1
    ssh -o ConnectTimeout=10 ubuntu@$vm_ip "hostname" 2>/dev/null || echo "unknown"
}

# Function to set hostname on a VM
set_vm_hostname() {
    local vm_ip=$1
    local new_hostname=$2
    local vm_desc=$3
    
    log "Setting hostname for $vm_desc ($vm_ip) to: $new_hostname"
    
    # Get current hostname
    local current_hostname=$(get_current_hostname $vm_ip)
    info "Current hostname: $current_hostname"
    
    # Set new hostname
    ssh -o ConnectTimeout=10 ubuntu@$vm_ip "sudo hostnamectl set-hostname $new_hostname" 2>/dev/null
    
    # Update /etc/hosts file
    ssh -o ConnectTimeout=10 ubuntu@$vm_ip "sudo sed -i 's/127.0.1.1.*/127.0.1.1\t$new_hostname/' /etc/hosts" 2>/dev/null
    
    # Verify the change
    local verify_hostname=$(get_current_hostname $vm_ip)
    if [ "$verify_hostname" = "$new_hostname" ]; then
        log "✓ Hostname successfully set to: $new_hostname"
    else
        error "Failed to set hostname. Current: $verify_hostname, Expected: $new_hostname"
        return 1
    fi
    
    echo
}

# Function to show current hostnames
show_current_hostnames() {
    echo
    echo "=========================================="
    echo "        Current VM Hostnames"
    echo "=========================================="
    echo
    
    echo "VM1 (Core Network) - $VM1:"
    echo "  Current: $(get_current_hostname $VM1)"
    echo "  Target:  $HOSTNAME_VM1"
    echo
    
    echo "VM2 (UPF) - $VM2:"
    echo "  Current: $(get_current_hostname $VM2)"
    echo "  Target:  $HOSTNAME_VM2"
    echo
    
    echo "VM3 (gNodeB) - $VM3:"
    echo "  Current: $(get_current_hostname $VM3)"
    echo "  Target:  $HOSTNAME_VM3"
    echo
    
    echo "VM4 (UE) - $VM4:"
    echo "  Current: $(get_current_hostname $VM4)"
    echo "  Target:  $HOSTNAME_VM4"
    echo
}

# Function to confirm changes
confirm_changes() {
    echo
    echo "=========================================="
    echo "        Hostname Change Summary"
    echo "=========================================="
    echo
    echo "The following hostname changes will be made:"
    echo
    echo "VM1 ($VM1): $(get_current_hostname $VM1) → $HOSTNAME_VM1"
    echo "VM2 ($VM2): $(get_current_hostname $VM2) → $HOSTNAME_VM2"
    echo "VM3 ($VM3): $(get_current_hostname $VM3) → $HOSTNAME_VM3"
    echo "VM4 ($VM4): $(get_current_hostname $VM4) → $HOSTNAME_VM4"
    echo
    echo "Note: These changes will take effect after reboot or re-login."
    echo
    read -p "Do you want to proceed with these changes? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        echo "Hostname changes cancelled."
        exit 0
    fi
}

# Function to test connectivity after changes
test_connectivity() {
    log "Testing connectivity with new hostnames..."
    
    local vms=("$VM1:$HOSTNAME_VM1" "$VM2:$HOSTNAME_VM2" "$VM3:$HOSTNAME_VM3" "$VM4:$HOSTNAME_VM4")
    
    for vm in "${vms[@]}"; do
        local ip=$(echo $vm | cut -d: -f1)
        local hostname=$(echo $vm | cut -d: -f2)
        
        echo -n "Testing $ip ($hostname)... "
        if ssh -o ConnectTimeout=5 ubuntu@$ip "hostname" 2>/dev/null | grep -q "$hostname"; then
            echo -e "${GREEN}✓${NC}"
        else
            echo -e "${RED}✗${NC}"
        fi
    done
}

# Function to show usage examples
show_examples() {
    echo
    echo "=========================================="
    echo "        New Hostname Usage Examples"
    echo "=========================================="
    echo
    echo "After reboot/re-login, you can use:"
    echo
    echo "SSH with hostnames:"
    echo "  ssh vm1          # Connect to vm1-5gc-core"
    echo "  ssh vm2          # Connect to vm2-5gc-upf"
    echo "  ssh vm3          # Connect to vm3-5gc-gnb"
    echo "  ssh vm4          # Connect to vm4-5gc-ue"
    echo
    echo "Check hostnames:"
    echo "  ssh vm1 'hostname'"
    echo "  ssh vm2 'hostname'"
    echo "  ssh vm3 'hostname'"
    echo "  ssh vm4 'hostname'"
    echo
    echo "=========================================="
}

# Main function
main() {
    echo
    echo "=========================================="
    echo "      VM Hostname Modification Script"
    echo "=========================================="
    echo "Date: $(date)"
    echo "=========================================="
    echo
    
    # Show current hostnames
    show_current_hostnames
    
    # Confirm changes
    confirm_changes
    
    log "Starting hostname modifications..."
    echo
    
    # Set hostnames for each VM
    set_vm_hostname $VM1 $HOSTNAME_VM1 "Core Network VM"
    set_vm_hostname $VM2 $HOSTNAME_VM2 "UPF VM"
    set_vm_hostname $VM3 $HOSTNAME_VM3 "gNodeB VM"
    set_vm_hostname $VM4 $HOSTNAME_VM4 "UE VM"
    
    echo
    log "Hostname modifications completed!"
    
    # Test connectivity
    test_connectivity
    
    # Show usage examples
    show_examples
    
    echo
    warn "IMPORTANT: Hostname changes will take full effect after reboot or re-login."
    echo "You may need to reboot the VMs or re-login to see the new hostnames."
    echo
    echo "To reboot all VMs:"
    echo "  ssh vm1 'sudo reboot'"
    echo "  ssh vm2 'sudo reboot'"
    echo "  ssh vm3 'sudo reboot'"
    echo "  ssh vm4 'sudo reboot'"
}

# Handle command line arguments
case "${1:-}" in
    "--help"|"-h")
        echo "Usage: $0 [OPTIONS]"
        echo
        echo "Options:"
        echo "  -h, --help          Show this help message"
        echo "  --show-current      Show current hostnames only"
        echo
        echo "This script will:"
        echo "  1. Show current hostnames of all VMs"
        echo "  2. Set new descriptive hostnames"
        echo "  3. Update /etc/hosts files"
        echo "  4. Test connectivity"
        echo
        echo "New hostnames will be:"
        echo "  VM1: vm1-5gc-core"
        echo "  VM2: vm2-5gc-upf"
        echo "  VM3: vm3-5gc-gnb"
        echo "  VM4: vm4-5gc-ue"
        ;;
    "--show-current")
        show_current_hostnames
        ;;
    *)
        main
        ;;
esac 